{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Validation notebook for ArcFace models\n",
    "\n",
    "## Overview\n",
    "Use this notebook to verify the accuracy of a trained ArcFace model in ONNX format on the validation datasets.\n",
    "\n",
    "## Models supported\n",
    "* LResNet100E-IR (ResNet100 backend with ArcFace loss)\n",
    "\n",
    "## Prerequisites\n",
    "The following packages need to be installed before proceeding:\n",
    "* Protobuf compiler - `sudo apt-get install protobuf-compiler libprotoc-dev` (required for ONNX. This will work for any linux system. For detailed installation guidelines head over to [ONNX documentation](https://github.com/onnx/onnx#installation))\n",
    "* ONNX - `pip install onnx`\n",
    "* MXNet - `pip install mxnet-cu90mkl --pre -U` (tested on this version GPU, can use other versions. `--pre` indicates a pre build of MXNet which is required here for ONNX version compatibility. `-U` uninstalls any existing MXNet version allowing for a clean install)\n",
    "* numpy - `pip install numpy`\n",
    "* matplotlib - `pip install matplotlib`\n",
    "* OpenCV - `pip install opencv-python`\n",
    "* Scikit-learn - `pip install scikit-learn`\n",
    "* EasyDict - `pip install easydict`\n",
    "\n",
    "In order to do inference with a python script:\n",
    "* Generate the script : In Jupyter Notebook browser, go to File -> Download as -> Python (.py)\n",
    "* Run the script: `python arcface_validation.py`"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Import dependencies\n",
    "Verify that all dependencies are installed using the cell below. Continue if no errors encountered, warnings can be ignored."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from __future__ import absolute_import\n",
    "from __future__ import division\n",
    "from __future__ import print_function\n",
    "\n",
    "import os\n",
    "import sys\n",
    "import numpy as np\n",
    "from scipy import misc\n",
    "from sklearn.model_selection import KFold\n",
    "from scipy import interpolate\n",
    "import sklearn\n",
    "import cv2\n",
    "import math\n",
    "import datetime\n",
    "import pickle\n",
    "import mxnet as mx\n",
    "from mxnet import ndarray as nd\n",
    "from easydict import EasyDict as edict\n",
    "from mxnet.contrib.onnx.onnx2mx.import_model import import_model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Data loading helper code\n",
    "* `load_bin()` loads validation datasets\n",
    "* `load_property()` loads information like image size, number of identities from the property file in the dataset folder"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_bin(path, image_size):\n",
    "    try:\n",
    "        # python 3\n",
    "        bins, issame_list = pickle.load(open(path, 'rb'), encoding='bytes')\n",
    "    except:\n",
    "        # python 2\n",
    "        bins, issame_list = pickle.load(open(path, 'rb'))\n",
    "    data_list = []\n",
    "    for flip in [0,1]:\n",
    "        data = nd.empty((len(issame_list)*2, 3, image_size[0], image_size[1]))\n",
    "        data_list.append(data)\n",
    "    for i in range(len(issame_list)*2):\n",
    "        _bin = bins[i]\n",
    "        img = mx.image.imdecode(_bin)\n",
    "        img = nd.transpose(img, axes=(2, 0, 1))\n",
    "        for flip in [0,1]:\n",
    "            if flip==1:\n",
    "                img = mx.ndarray.flip(data=img, axis=2)\n",
    "            data_list[flip][i][:] = img\n",
    "        if i%1000==0:\n",
    "            print('loading bin', i)\n",
    "    print(data_list[0].shape)\n",
    "    return (data_list, issame_list)\n",
    "\n",
    "def load_property(data_dir):\n",
    "    prop = edict()\n",
    "    for line in open(os.path.join(data_dir, 'property')):\n",
    "        vec = line.strip().split(',')\n",
    "        assert len(vec)==3\n",
    "        prop.num_classes = int(vec[0])\n",
    "        prop.image_size = [int(vec[1]), int(vec[2])]\n",
    "    return prop"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Evaluation helper code\n",
    "* `class LFold` is used to split the data for K-fold crossvalidation\n",
    "* `calculate_roc()` computes accuracy on each fold of cross validation using `calculate_accuracy()`\n",
    "* `calculate_accuracy()` computes the actual accuracy for test samples by thresholding the distance between embedding vectors of a test image pair and comparing the output with the ground truth for the pair\n",
    "* `evaluate()` splits embeddings into test pairs and computes accuracies using `calculate_roc()`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "class LFold:\n",
    "    def __init__(self, n_splits = 2, shuffle = False):\n",
    "        self.n_splits = n_splits\n",
    "        if self.n_splits>1:\n",
    "            self.k_fold = KFold(n_splits = n_splits, shuffle = shuffle)\n",
    "\n",
    "    def split(self, indices):\n",
    "        if self.n_splits>1:\n",
    "            return self.k_fold.split(indices)\n",
    "        else:\n",
    "            return [(indices, indices)]\n",
    "\n",
    "\n",
    "def calculate_roc(thresholds, embeddings1, embeddings2, actual_issame, nrof_folds=10):\n",
    "    assert(embeddings1.shape[0] == embeddings2.shape[0])\n",
    "    assert(embeddings1.shape[1] == embeddings2.shape[1])\n",
    "    nrof_pairs = min(len(actual_issame), embeddings1.shape[0])\n",
    "    nrof_thresholds = len(thresholds)\n",
    "    k_fold = LFold(n_splits=nrof_folds, shuffle=False)\n",
    "    \n",
    "    tprs = np.zeros((nrof_folds,nrof_thresholds))\n",
    "    fprs = np.zeros((nrof_folds,nrof_thresholds))\n",
    "    accuracy = np.zeros((nrof_folds))\n",
    "    indices = np.arange(nrof_pairs)\n",
    "    \n",
    "    diff = np.subtract(embeddings1, embeddings2)\n",
    "    dist = np.sum(np.square(diff),1)\n",
    "    \n",
    "    for fold_idx, (train_set, test_set) in enumerate(k_fold.split(indices)):        \n",
    "        # Find the best threshold for the fold\n",
    "        acc_train = np.zeros((nrof_thresholds))\n",
    "        for threshold_idx, threshold in enumerate(thresholds):\n",
    "            _, _, acc_train[threshold_idx] = calculate_accuracy(threshold, dist[train_set], actual_issame[train_set])\n",
    "        best_threshold_index = np.argmax(acc_train)\n",
    "        for threshold_idx, threshold in enumerate(thresholds):\n",
    "            tprs[fold_idx,threshold_idx], fprs[fold_idx,threshold_idx], _ = calculate_accuracy(threshold, dist[test_set], actual_issame[test_set])\n",
    "        _, _, accuracy[fold_idx] = calculate_accuracy(thresholds[best_threshold_index], dist[test_set], actual_issame[test_set])\n",
    "          \n",
    "    tpr = np.mean(tprs,0)\n",
    "    fpr = np.mean(fprs,0)\n",
    "    return tpr, fpr, accuracy\n",
    "\n",
    "def calculate_accuracy(threshold, dist, actual_issame):\n",
    "    predict_issame = np.less(dist, threshold)\n",
    "    tp = np.sum(np.logical_and(predict_issame, actual_issame))\n",
    "    fp = np.sum(np.logical_and(predict_issame, np.logical_not(actual_issame)))\n",
    "    tn = np.sum(np.logical_and(np.logical_not(predict_issame), np.logical_not(actual_issame)))\n",
    "    fn = np.sum(np.logical_and(np.logical_not(predict_issame), actual_issame))\n",
    "  \n",
    "    tpr = 0 if (tp+fn==0) else float(tp) / float(tp+fn)\n",
    "    fpr = 0 if (fp+tn==0) else float(fp) / float(fp+tn)\n",
    "    acc = float(tp+tn)/dist.size\n",
    "    return tpr, fpr, acc\n",
    "\n",
    "def evaluate(embeddings, actual_issame, nrof_folds=10):\n",
    "    # Calculate evaluation metrics\n",
    "    thresholds = np.arange(0, 4, 0.01)\n",
    "    embeddings1 = embeddings[0::2]\n",
    "    embeddings2 = embeddings[1::2]\n",
    "    tpr, fpr, accuracy = calculate_roc(thresholds, embeddings1, embeddings2,\n",
    "        np.asarray(actual_issame), nrof_folds=nrof_folds)\n",
    "    return accuracy"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Define test function\n",
    "`test()` takes a validation set `data_set` and the MXNet model `mx_model` as input and computes accuracies using `evaluate()` on the set using `nfolds` cross validation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def test(data_set, mx_model, batch_size, nfolds=10, data_extra = None, label_shape = None):\n",
    "    print('testing verification..')\n",
    "    data_list = data_set[0]\n",
    "    issame_list = data_set[1]\n",
    "    model = mx_model\n",
    "    embeddings_list = []\n",
    "    if data_extra is not None:\n",
    "        _data_extra = nd.array(data_extra)\n",
    "    time_consumed = 0.0\n",
    "    if label_shape is None:\n",
    "        _label = nd.ones( (batch_size,) )\n",
    "    else:\n",
    "        _label = nd.ones( label_shape )\n",
    "    for i in range( len(data_list) ):\n",
    "        data = data_list[i]\n",
    "        embeddings = None\n",
    "        ba = 0\n",
    "        while ba<data.shape[0]:\n",
    "            bb = min(ba+batch_size, data.shape[0])\n",
    "            count = bb-ba\n",
    "            _data = nd.slice_axis(data, axis=0, begin=bb-batch_size, end=bb)\n",
    "            time0 = datetime.datetime.now()\n",
    "            if data_extra is None:\n",
    "                db = mx.io.DataBatch(data=(_data,), label=(_label,))\n",
    "            else:\n",
    "                db = mx.io.DataBatch(data=(_data,_data_extra), label=(_label,))\n",
    "            model.forward(db, is_train=False)\n",
    "            net_out = model.get_outputs()\n",
    "            _embeddings = net_out[0].asnumpy()\n",
    "            time_now = datetime.datetime.now()\n",
    "            diff = time_now - time0\n",
    "            time_consumed+=diff.total_seconds()\n",
    "            if embeddings is None:\n",
    "                embeddings = np.zeros( (data.shape[0], _embeddings.shape[1]) )\n",
    "            embeddings[ba:bb,:] = _embeddings[(batch_size-count):,:]\n",
    "            ba = bb\n",
    "        embeddings_list.append(embeddings)\n",
    "\n",
    "    _xnorm = 0.0\n",
    "    _xnorm_cnt = 0\n",
    "    for embed in embeddings_list:\n",
    "        for i in range(embed.shape[0]):\n",
    "            _em = embed[i]\n",
    "            _norm=np.linalg.norm(_em)\n",
    "            _xnorm+=_norm\n",
    "            _xnorm_cnt+=1\n",
    "    _xnorm /= _xnorm_cnt\n",
    "\n",
    "    embeddings = embeddings_list[0].copy()\n",
    "    embeddings = sklearn.preprocessing.normalize(embeddings)\n",
    "    embeddings = embeddings_list[0] + embeddings_list[1]\n",
    "    embeddings = sklearn.preprocessing.normalize(embeddings)\n",
    "    print(embeddings.shape)\n",
    "    print('infer time', time_consumed)\n",
    "    accuracy = evaluate(embeddings, issame_list, nrof_folds=nfolds)\n",
    "    acc2, std2 = np.mean(accuracy), np.std(accuracy)\n",
    "    return acc2, std2, _xnorm, embeddings_list"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Specify paths and parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Path to dataset\n",
    "data_dir = '/home/ubuntu/faces_ms1m_112x112/'\n",
    "# Path to model file\n",
    "model = '/home/ubuntu/resnet100.onnx'\n",
    "# Verification targets\n",
    "target = 'lfw,cfp_ff,cfp_fp,agedb_30'\n",
    "# Batch size\n",
    "batch_size = 64\n",
    "# Number of folds for cross validation\n",
    "nfolds = 10"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Compute validation accuracies\n",
    "This is the main function and it loads the ONNX model, loads the validation datasets and computes the validation accuracies on the data according to the verification targets specified above. The accuracy for each target is displayed in the output of the cell."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "image_size [112, 112]\n",
      "model loading time 28.645658\n",
      "loading..  lfw\n",
      "loading bin 0\n",
      "loading bin 1000\n",
      "loading bin 2000\n",
      "loading bin 3000\n",
      "loading bin 4000\n",
      "loading bin 5000\n",
      "loading bin 6000\n",
      "loading bin 7000\n",
      "loading bin 8000\n",
      "loading bin 9000\n",
      "loading bin 10000\n",
      "loading bin 11000\n",
      "(12000, 3, 112, 112)\n",
      "loading..  cfp_ff\n",
      "loading bin 0\n",
      "loading bin 1000\n",
      "loading bin 2000\n",
      "loading bin 3000\n",
      "loading bin 4000\n",
      "loading bin 5000\n",
      "loading bin 6000\n",
      "loading bin 7000\n",
      "loading bin 8000\n",
      "loading bin 9000\n",
      "loading bin 10000\n",
      "loading bin 11000\n",
      "loading bin 12000\n",
      "loading bin 13000\n",
      "(14000, 3, 112, 112)\n",
      "loading..  cfp_fp\n",
      "loading bin 0\n",
      "loading bin 1000\n",
      "loading bin 2000\n",
      "loading bin 3000\n",
      "loading bin 4000\n",
      "loading bin 5000\n",
      "loading bin 6000\n",
      "loading bin 7000\n",
      "loading bin 8000\n",
      "loading bin 9000\n",
      "loading bin 10000\n",
      "loading bin 11000\n",
      "loading bin 12000\n",
      "loading bin 13000\n",
      "(14000, 3, 112, 112)\n",
      "loading..  agedb_30\n",
      "loading bin 0\n",
      "loading bin 1000\n",
      "loading bin 2000\n",
      "loading bin 3000\n",
      "loading bin 4000\n",
      "loading bin 5000\n",
      "loading bin 6000\n",
      "loading bin 7000\n",
      "loading bin 8000\n",
      "loading bin 9000\n",
      "loading bin 10000\n",
      "loading bin 11000\n",
      "(12000, 3, 112, 112)\n",
      "testing verification..\n",
      "(12000, 512)\n",
      "infer time 41.15619399999997\n",
      "[lfw]XNorm: 22.453576\n",
      "[lfw]Accuracy-Flip: 0.99767+-0.00271\n",
      "testing verification..\n",
      "(14000, 512)\n",
      "infer time 46.53032800000003\n",
      "[cfp_ff]XNorm: 21.054120\n",
      "[cfp_ff]Accuracy-Flip: 0.99829+-0.00270\n",
      "testing verification..\n",
      "(14000, 512)\n",
      "infer time 46.556721\n",
      "[cfp_fp]XNorm: 21.643374\n",
      "[cfp_fp]Accuracy-Flip: 0.94214+-0.01027\n",
      "testing verification..\n",
      "(12000, 512)\n",
      "infer time 39.959474000000014\n",
      "[agedb_30]XNorm: 23.488695\n",
      "[agedb_30]Accuracy-Flip: 0.97817+-0.00765\n"
     ]
    }
   ],
   "source": [
    "if __name__ == '__main__':\n",
    "    \n",
    "    # Load image size\n",
    "    prop = load_property(data_dir)\n",
    "    image_size = prop.image_size\n",
    "    print('image_size', image_size)\n",
    "    \n",
    "    # Determine and set context\n",
    "    if len(mx.test_utils.list_gpus())==0:\n",
    "        ctx = mx.cpu()\n",
    "        batch_size=1\n",
    "    else:\n",
    "        ctx = mx.gpu(0)\n",
    "        \n",
    "    time0 = datetime.datetime.now()\n",
    "    \n",
    "    # Import ONNX model\n",
    "    sym, arg_params, aux_params = import_model(model)\n",
    "    all_layers = sym.get_internals()\n",
    "    sym = all_layers['fc1_output']\n",
    "    # Define model\n",
    "    model = mx.mod.Module(symbol=sym, context=ctx, label_names = None)\n",
    "    # Bind parameters to the model\n",
    "    model.bind(data_shapes=[('data', (batch_size, 3, image_size[0], image_size[1]))])\n",
    "    model.set_params(arg_params, aux_params)\n",
    "    time_now = datetime.datetime.now()\n",
    "    diff = time_now - time0\n",
    "    print('model loading time', diff.total_seconds())\n",
    "\n",
    "    ver_list = []\n",
    "    ver_name_list = []\n",
    "    \n",
    "    # Iterate over verification targets\n",
    "    for name in target.split(','):\n",
    "        path = os.path.join(data_dir,name+\".bin\")\n",
    "        # Load data\n",
    "        if os.path.exists(path):\n",
    "            print('loading.. ', name)\n",
    "            data_set = load_bin(path, image_size)\n",
    "            ver_list.append(data_set)\n",
    "            ver_name_list.append(name)\n",
    "    \n",
    "    # Iterate over verification targets\n",
    "    for i in range(len(ver_list)):\n",
    "        # Compute and print validation accuracies\n",
    "        acc2, std2, xnorm, embeddings_list = test(ver_list[i], model, batch_size, nfolds)\n",
    "        print('[%s]XNorm: %f' % (ver_name_list[i], xnorm))\n",
    "        print('[%s]Accuracy-Flip: %1.5f+-%1.5f' % (ver_name_list[i], acc2, std2))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Environment (conda_face3)",
   "language": "python",
   "name": "conda_face3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
